from vector import Vector

class Matrix:

    def __init__(self, lst2d):
        self._values = list(lst2d)
        self._m = len(self._values)
        self._n = len(self._values[0])


    def __repr__(self):
        return f"Matrix({self._values})"

    __str__ = __repr__

    def row_num(self):
        return self._m

    def col_num(self):
        return self._n

    def shape(self):
        return self._m,self._n

    def row_vector(self,index):
        return Vector(self._values[index])

    def col_vector(self,index):
        return Vector(e[index] for e in self._values)

    def __add__(self, other):
        assert self.shape() == other.shape(), \
            "Error in add,Vector shape must be same"
        return Matrix([[self._values[i][j] + other.Ab[i][j] for j in range(self.col_num())] for i in range(self.row_num())])

    def __sub__(self, other):
        assert self.shape() == other.shape(), \
            "Error in sub,Vector shape must be same"
        return Matrix([[a-b for a,b in zip(self.row_vector(i),other.row_vector(i))] for i in range(self.row_num())])

    def __mul__(self, k):
        return Matrix([[e*k for e in self.row_vector(i)] for i in range(self.row_num())])

    def __rmul__(self, k):
        return Matrix([[e*k for e in self.row_vector(i)] for i in range(self.row_num())])

    def __truediv__(self, k):
        return self * (1/k)

    def __pos__(self):
        return self * 1

    def __neg__(self):
        return self * -1

    @classmethod
    def zeros(cls,r,c):
        return cls([[0]*c for _ in range(r)])

    @classmethod
    def ones(cls,r,c):
        return cls([[1]*c for _ in range(r)])

    def dot(self,other):

        if isinstance(other,Vector):
            assert self.col_num() == len(other), \
            "错误:Matrix.dot(Matrix,Vector) 矩阵的col_num 必须等于向量的大小"
            return Vector( [ self.row_vector(i).dot(other) for i in range(self.row_num())])
        if isinstance(other,Matrix):
            assert self.col_num() == other.row_num(), \
                "错误:Matrix.dot(Matrix,Matrix) 矩阵的col_num 必须等于other的row_num"
            return Matrix([ [ self.row_vector(i).dot(other.col_vector(j)) for j in range(other.col_num())]for i in range(self.row_num())])

    def T(self):
        m = self._values
        for i in range(self.row_num()):
            for j in range(self.col_num()):
                if i<j:
                    m[i][j],m[j][i] = m[j][i],m[i][j]
        return Matrix(m)


if __name__ == '__main__':
    m1 = Matrix([[1,2],[3,4]])
    print(type(m1),m1)
    print(m1._m,m1._n)

    print("--------------")
    print(type(m1.row_vector(0)))
    print(m1.row_vector(1))
    print(m1.col_vector(1))

    print("--------------")
    m1 = Matrix([[1, 2,3], [4, 5, 6],[7,8,9]])
    m2 = Matrix([[11, 12, 13], [14, 15, 16], [17, 18, 19]])

    print(m1+m2)
    print(m1 - m2)

    print("--------------")
    print(m1*10)
    print(10*m1)
    print(m1/10)

    print("--------------")
    print(Matrix.zeros(4,5))

    print("--------------")
    m1 = Matrix([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
    v1 = Vector([1,2,3])
    print(m1.dot(v1))
    print(m1.dot(m1))

    print("--------------")
    print(m1.T())